#pragma once
#include <boost/beast/core.hpp>
#include <boost/beast/websocket.hpp>
#include <boost/beast/ssl.hpp>
#include <boost/asio/connect.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <boost/asio/ssl.hpp>
#include <functional>
#include <memory>
#include <string>
#include <iostream>

class WSConnection : public std::enable_shared_from_this<WSConnection> {
public:
    using MessageHandler = std::function<void(const std::string&)>;
    using ErrorHandler = std::function<void(const std::string&)>;
    using ConnectHandler = std::function<void()>;

    WSConnection(boost::asio::io_context& ioc, boost::asio::ssl::context& ssl_ctx);
    ~WSConnection();

    void connect(const std::string& host, const std::string& port, const std::string& path);
    void close();
    void send(const std::string& message);
    bool isOpen() const;

    void setMessageHandler(MessageHandler handler) { onMessage_ = std::move(handler); }
    void setErrorHandler(ErrorHandler handler) { onError_ = std::move(handler); }
    void setConnectHandler(ConnectHandler handler) { onConnect_ = std::move(handler); }

private:
    void onResolve(const boost::system::error_code& ec,
                   boost::asio::ip::tcp::resolver::results_type results);
    void onConnect(const boost::system::error_code& ec);
    void onSSLHandshake(const boost::system::error_code& ec);
    void onWSHandshake(const boost::system::error_code& ec);
    void doRead();
    void onRead(const boost::system::error_code& ec, std::size_t bytes_transferred);
    void onWrite(const boost::system::error_code& ec, std::size_t bytes_transferred);
    void onClose(const boost::system::error_code& ec);

private:
    boost::asio::io_context& ioc_;
    boost::asio::ssl::context& ssl_ctx_;
    boost::asio::ip::tcp::resolver resolver_;

    std::unique_ptr<boost::beast::websocket::stream<
        boost::beast::ssl_stream<boost::asio::ip::tcp::socket>>> ws_;

    boost::beast::flat_buffer buffer_;
    std::string host_;
    std::string path_;

    MessageHandler onMessage_;
    ErrorHandler onError_;
    ConnectHandler onConnect_;
};
